#ifndef WEBSOCKET_SERVER_H
#define WEBSOCKET_SERVER_H

#include <thread>
#include <vector>
#include <memory>
#include <atomic>
#include <functional>
#include <unordered_map>
#include <mutex>
#include <string>
#include <sys/socket.h>
#include <netinet/in.h>

namespace websocket {

class WebSocketSession;

class WebSocketServer {
public:
    using MessageHandler = std::function<void(std::shared_ptr<WebSocketSession>, const std::string&)>;
    using ConnectionHandler = std::function<void(std::shared_ptr<WebSocketSession>)>;
    using DisconnectionHandler = std::function<void(std::shared_ptr<WebSocketSession>)>;

    WebSocketServer(int port = 8080);
    ~WebSocketServer();
    
    // 启动和停止服务器
    bool start();
    void stop();
    
    // 设置回调函数
    void setMessageHandler(MessageHandler handler) { message_handler_ = handler; }
    void setConnectionHandler(ConnectionHandler handler) { connection_handler_ = handler; }
    void setDisconnectionHandler(DisconnectionHandler handler) { disconnection_handler_ = handler; }
    
    // 广播消息给所有客户端
    void broadcast(const std::string& message);

    // 广播消息给指定 socket 列表的客户端
    void broadcastTo(const std::vector<int>& sockets, const std::string& message);
    
    // 获取连接统计信息
    size_t getConnectionCount() const;

private:
    void acceptConnections();
    void handleClient(int client_socket);
    bool performWebSocketHandshake(int client_socket, const std::string& request);
    std::string generateWebSocketAcceptKey(const std::string& client_key);
    bool sendData(int socket, const char* data, size_t length);
    
    int port_;
    int server_socket_;
    std::atomic<bool> running_;
    std::thread accept_thread_;
    
    MessageHandler message_handler_;
    ConnectionHandler connection_handler_;
    DisconnectionHandler disconnection_handler_;
    
    std::unordered_map<int, std::shared_ptr<WebSocketSession>> sessions_;
    mutable std::mutex sessions_mutex_;
};

} // namespace websocket

#endif // WEBSOCKET_SERVER_H